# -*- coding: utf-8 -*-
from tlib import yTelnet

def config_manager():
    def _deco(func):
        def _func(self, *args, **kwargs):
            # 注意 self 是作为参数传进来的
            self.change_to_config_mode()
            func(self, *args, **kwargs)
            self.change_to_enable_mode()
        return _func
    return _deco

class deviceControl(yTelnet.TelnetSock):
    '''
           此类用于控制设备配置，为不同设备不同的配置方式提供统一的接口 。
    dut_type must in ['T1-ADC', 'CISCO', 'T1-FW', 'HUAWEI']   ---HUAWEI暂不支持
    cisco交换机的enable密码请配置成cisco
    关于配置的注意事项：
        涉及到接口的配置，每次都会清除掉接口之前的配置，防止前面配置引起错误，所以配置之间不要有依赖；
    '''
    def __init__(self, dut_type, dutname, destip, dport = 23, username = 'admin', 
                 password = 'adc.admin', log_level = 'warn'):
        self.dut_type = dut_type
        yTelnet.TelnetSock.__init__(self, destip, dport, 
                                    dutname, username, password, log_level = log_level, sendInterval=0)
        self.login_enable()
    def login_enable(self):
        if self.dut_type == 'CISCO':
            self.post('enable')
            self.recv_until('Password:')
            self.post('cisco')
        self.send_commands('enable\nter idl 0')
    def change_to_config_mode(self):
        self.post('conf t')
    def change_to_enable_mode(self):
        self.post('end')
    @config_manager()
    def create_vlan(self, vlan_id):
        self.post('vlan %s'%vlan_id)
    @config_manager()
    def del_vlan(self, vlan_id):
        self.post('no vlan %s'%vlan_id)
    @config_manager()
    def untag_vlan(self, interface_name, vlan_id):
        self.default_interface_config(interface_name)
        self.post('interface %s'%interface_name)
        if self.dut_type == 'CISCO':
            self.send_commands('sw\nsw mo access\nsw access vlan %s'%vlan_id)
        else: 
            self.post('untag vlan %s'%vlan_id)
    @config_manager()
    def tag_vlan(self, interface_name, vlan_id):
        self.default_interface_config(interface_name)
        self.post('int %s'%interface_name)
        if self.dut_type == 'CISCO':
            self.send_commands('sw\nsw trunk en do\nsw mo trunk\nsw tr allow vlan %s'%vlan_id)
        else: 
            self.post('tag vlan %s'%vlan_id)
    @config_manager()
    def add_ip(self, interface_name, ip, mask):
        "mask must be int smaller than 32"
        self.post('interface %s'%interface_name)
        if self.dut_type == 'CISCO':
            self.post('ip addr %s %s'%(ip, self._mask_format(mask)))
        else:
            self.post('ip addr %s/%s'%(ip, mask))
            self.post('allow ping')
    @config_manager()
    def add_ip6(self, interface_name, ip, mask):
        "mask must be int smaller than 32"
        self.post('interface %s'%interface_name)
        self.post('ipv6 addr %s/%s'%(ip, mask))
        self.post('allow ping')
    def _mask_format(self, mask):
        res = []
        a, b = divmod(mask, 8)
        for i in xrange(a):
            res.append('255')
        if b:
            res.append(str(int((b*'1' + (8-b)*'0'), 2)))
        if len(res)<4:
            res.append(((4-len(res))*'.0').strip('.'))
        return '.'.join(res)
    @config_manager()
    def create_trunk(self, trunk_id, interfaces, mode = 'static'):
        '''interface mask be list
        trunk_id is int'''
        if self.dut_type == 'CISCO':
            pass
        else:
            self.post('interface tvi%s'%trunk_id)
            self.default_interface_config('tvi%s'%trunk_id)
            for i in interfaces:
                self.default_interface_config(i)
                self.post('interface %s'%i)
                self.post('trunk-group %s'%trunk_id)
    @config_manager()
    def default_interface_config(self, interface_name):
        '''清空某接口的配置'''
        if self.dut_type == 'CISCO':
            self.post('default interface %s'%interface_name)
        else:
            a = self.post('show run interface')
            b = a[a.find('interface %s'%interface_name):]
            if not b:
                self.logger.warn('WARNING: interface %s is not exist!')
                return
            c = b[0:b.find('!')]
            cl = [i.strip() for i in c.split('\r\n') if i]
            if cl:
                for l in cl:
                    if l.find('--More--')!=-1:
                        cl.remove(l)
                self.post(cl[0])
                for i in cl[1:]:
                    if i.find('trunk-group')!=-1:
                        self.post('no trunk-group')
                    elif i.find('float')!=-1:
                        self.post('no %s'%i[:i.find('unit')])
                    elif i.find('fc on')!=-1:
                        pass
                    else:
                        self.post('no %s'%i)
            else:
                self.logger.warn('WARNING: interface %s config is none'%interface_name)
    def exit(self):
        try:
            self.send_commands('end\nexit\n')
            self._recv()
            self.close()
            self.logger.exit()
        except:
            #print "Maybe socket was dead!"
            return

def interface_config_handle(command, dc, interface_name, vlan_id = 0, trunk_id = 0):
    if command == 'tag vlan':
        dc.tag_vlan(interface_name, vlan_id)
    elif command == 'untag vlan':
        dc.untag_vlan(interface_name, vlan_id)
    elif command == 'trunk tag vlan':
        if dc.dut_type == 'CISCO':
            dc.tag_vlan(interface_name, vlan_id)
        else:
            dc.create_trunk(trunk_id, [interface_name, ])
            dc.tag_vlan('tvi%s'%trunk_id, vlan_id)
    elif command == 'trunk untag vlan':
        if dc.dut_type == 'CISCO':
            dc.untag_vlan(interface_name, vlan_id)
        else:
            try:
                dc.create_trunk(trunk_id, [interface_name, ])
                dc.untag_vlan('tvi%s'%trunk_id, vlan_id)
            except:
                print "%s, %s error!!!"%(trunk_id, interface_name)
    elif command == 'trunk':
        dc.post('interface vlan%s'%vlan_id)
        dc.post('no ip addr')
        dc.post('no ipv6 add')
        dc.create_trunk(trunk_id, [interface_name, ])
    elif command == 'physic':
        dc.post('interface vlan%s'%vlan_id)
        dc.post('no ip addr')
        dc.post('no ipv6 add')
        dc.default_interface_config(interface_name)
    else:
        pass

def add_ip_addr(command, dc, interface_name, ip, mask, ip6, prefix, vlan_id = 0, trunk_id = 0):
    if command.find('vlan')!=-1:
        dc.add_ip('vlan%s'%vlan_id, ip, mask)
        dc.add_ip6('vlan%s'%vlan_id, ip6, prefix)
    elif command == 'trunk':
        if dc.dut_type == 'CISCO':
            dc.post('interface %s'%interface_name)
            dc.post('no sw')
            dc.add_ip(interface_name, ip, mask)
            dc.add_ip6(interface_name, ip6, prefix)
        else:
            dc.add_ip('tvi%s'%trunk_id, ip, mask)
            dc.add_ip6('tvi%s'%trunk_id, ip6, prefix)
    elif command == 'physic':
        if dc.dut_type == 'CISCO':
            dc.post('interface %s'%interface_name)
            dc.post('no sw')
        else:
            pass
        dc.add_ip(interface_name, ip, mask)
        dc.add_ip6(interface_name, ip6, prefix)
    else:
        pass



def traffic(object):
    pass
    
    
if __name__ == '__main__':
    cisco = deviceControl('T1-FW' ,'yb_2u', '172.16.1.38', username = 'admin', password = 'fw.admin', log_level='info')
    cisco.ifPrint = 1
    cisco.exit()